from .base import BaseLLM
from openai import OpenAI
from config import DOUBAO_API_KEY, DOUBAO_BASE_URL, DOUBAO_MODEL


class DoubaoLLM(BaseLLM):
    def __init__(self):
        self.client = OpenAI(base_url=DOUBAO_BASE_URL, api_key=DOUBAO_API_KEY)

    def generate(self, prompt, **kwargs):
        resp = self.client.chat.completions.create(
            model=DOUBAO_MODEL,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": prompt,
                        },
                    ],
                }
            ],
            temperature=kwargs.get("temperature", 0.3),
            max_tokens=kwargs.get("max_tokens", 4096),
        )
        return resp.choices[0].message.content
